#!/usr/bin/env python
# coding: utf-8

# 导入必要的库
import numpy as np  # 导入 numpy 用于线性代数运算
import pandas as pd  # 导入 pandas 用于数据处理，例如读取 CSV 文件

# 从 dataclasses 导入 dataclass 用于创建数据类
from dataclasses import dataclass

# 从 peft 库导入 PeftModel 用于加载微调后的模型
from peft import PeftModel

# 从 transformers 库导入 AutoModelForCausalLM 用于加载因果语言模型，AutoTokenizer 用于加载分词器
from transformers import AutoModelForCausalLM, AutoTokenizer

# 导入 polars 用于快速数据处理
import polars as pl
import time  # 导入 time 用于时间相关操作
import torch  # 导入 torch 用于深度学习


# 配置类
@dataclass
class Config:
    qwen_dir = "/kaggle/input/qwen2.5/transformers/1.5b-instruct/1"  # Qwen 模型目录
    max_length = 2048  # 最大输入长度
    device = torch.device("cuda")  # 使用 CUDA 设备进行计算


cfg = Config()


# 加载数据的函数
def load_data(file_path):
    return pl.read_parquet(
        file_path
    ).to_pandas()  # 读取 parquet 文件并转换为 pandas DataFrame


# 预处理数据的函数
def preprocess_data(df):
    df.prompt = df.prompt.fillna("")  # 将 prompt 列中的缺失值填充为空字符串
    df.response_a = df.response_a.fillna("")  # 将 response_a 列中的缺失值填充为空字符串
    df.response_b = df.response_b.fillna("")  # 将 response_b 列中的缺失值填充为空字符串
    return df


# 初始化模型和分词器
tokenizer = AutoTokenizer.from_pretrained(cfg.qwen_dir)  # 从预训练模型目录加载分词器
model = AutoModelForCausalLM.from_pretrained(
    cfg.qwen_dir
)  # 从预训练模型目录加载因果语言模型
model.eval()  # 将模型设置为评估模式

# 将数据标记化的函数
instruction = """在下面为您提供的文本中，PROMPT是提出的问题；MODEL_A是第一个模型的回答；MODEL_B是第二个模型的回答。请从提供的两个回答中选择最佳答案。如果第一个答案更好，返回“model_a”；如果第二个答案更好，返回“model_b”。"""


def tokenize_data(df):
    tokenised_data = []  # 存储标记化后的数据
    ids = []  # 存储数据 ID
    for idx in range(len(df)):
        rec = df.iloc[idx, :]  # 获取当前行
        prompt = "PROMPT: " + rec["prompt"]  # 构建 prompt 文本
        model_a = "MODEL_A: " + rec["response_a"]  # 构建 model_a 文本
        model_b = "MODEL_B: " + rec["response_b"]  # 构建 model_b 文本
        text = prompt + model_a + model_b  # 合并文本
        messages = [
            {"role": "system", "content": instruction},  # 添加系统指令
            {"role": "user", "content": text},  # 添加用户输入
        ]
        text = tokenizer.apply_chat_template(
            messages,  # 将消息转换为模型输入格式
            tokenize=False,  # 不进行分词
            add_generation_prompt=True,  # 添加生成提示
        )

        tokenised_datum = tokenizer(
            [text],  # 输入文本
            return_tensors="pt",  # 返回 PyTorch 张量
            padding=True,  # 进行填充
            truncation=True,  # 进行截断
        ).to(
            cfg.device
        )  # 将张量移动到指定设备
        tokenised_data.append(tokenised_datum)  # 将标记化后的数据添加到列表
        ids.append(rec["id"])  # 将 ID 添加到列表
    return tokenised_data, ids


# 进行推理的函数
def inference(tokenised_datum):
    model.to(cfg.device)  # 将模型移动到指定设备
    input_ids = tokenised_datum.input_ids  # 获取输入 ID
    generated_ids = model.generate(
        input_ids, max_new_tokens=5  # 输入 ID  # 最大生成 token 数量
    )
    generated_ids = [
        output_ids[len(input_ids) :]
        for input_ids, output_ids in zip(input_ids, generated_ids)  # 获取生成的 ID
    ]

    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[
        0
    ]  # 将生成的 ID 解码为文本
    return response


# 设置环境变量
import os

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = (
    "expandable_segments:True"  # 允许 CUDA 内存动态扩展
)
# 加载训练数据
train = load_data("/kaggle/input/wsdm-cup-multilingual-chatbot-arena/train.parquet")
train = train.head(25000)  # 取前 25000 行数据
train = preprocess_data(train)  # 预处理训练数据
tokenised_data_train, ids = tokenize_data(train)  # 标记化训练数据
predictions_train = []  # 存储训练数据的预测结果
for tokenised_datum_train in tokenised_data_train:
    try:
        # torch.cuda.empty_cache()  # 清空 GPU 缓存
        response = inference(tokenised_datum_train)  # 进行推理
        predictions_train.append(response)  # 将预测结果添加到列表
    except Exception as e:
        print(f"An error occurred: {e}")  # 打印错误信息
        predictions_train.append("model_a")  # 如果出现错误，则预测为 "model_a"

# 加载测试数据
test = load_data("/kaggle/input/wsdm-cup-multilingual-chatbot-arena/test.parquet")
test = preprocess_data(test)  # 预处理测试数据
tokenised_data, ids = tokenize_data(test)  # 标记化测试数据
predictions = []  # 存储测试数据的预测结果
for tokenised_datum in tokenised_data:
    try:
        response = inference(tokenised_datum)  # 进行推理
        predictions.append(response)  # 将预测结果添加到列表
    except Exception as e:
        print(f"An error occurred: {e}")  # 打印错误信息
        predictions.append("model_a")  # 如果出现错误，则预测为 "model_a"

# 创建提交文件
submission = pd.DataFrame({"id": ids, "winner": predictions})  # ID 列  # 预测结果列
submission.to_csv("submission.csv", index=False)  # 将结果保存为 CSV 文件
print(submission)  # 打印提交结果
